package com.wave.wavesystem.ai.advisor;

import com.wave.wavesystem.ai.pojo.entity.AiConversation;
import com.wave.wavesystem.ai.service.IAiConversationService;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClientRequest;
import org.springframework.ai.chat.client.ChatClientResponse;
import org.springframework.ai.chat.client.advisor.api.CallAdvisor;
import org.springframework.ai.chat.client.advisor.api.CallAdvisorChain;
import org.springframework.ai.chat.client.advisor.api.StreamAdvisor;
import org.springframework.ai.chat.client.advisor.api.StreamAdvisorChain;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;

/**
 * token 计算
 */
@Slf4j
@Component
public class WaveTokenCalculateAdvisor implements CallAdvisor, StreamAdvisor {

    @Resource
    private IAiConversationService conversationService;

    @Override
    public ChatClientResponse adviseCall(ChatClientRequest chatClientRequest, CallAdvisorChain callAdvisorChain) {

        beforeRequest(chatClientRequest);

        ChatClientResponse chatClientResponse = callAdvisorChain.nextCall(chatClientRequest);

        log.info("adviseCall chatClientResponse={}", chatClientResponse);
        afterCallRequest(chatClientResponse);

        return chatClientResponse;
    }

    private void afterCallRequest(ChatClientResponse chatClientResponse) {
        Integer totalTokens = chatClientResponse.chatResponse().getMetadata().getUsage().getTotalTokens();

        String tokens = String.valueOf(totalTokens.intValue());
        String conversationId = (String) chatClientResponse.context().get(ChatMemory.CONVERSATION_ID);
        conversationService.updateChain()
                .eq(AiConversation::getConversationId, conversationId)
                .setRaw(AiConversation::getTotalTokens, "totalTokens + " + tokens)
                .update();

    }

    private void beforeRequest(ChatClientRequest chatClientRequest) {
    }


    @Override
    public Flux<ChatClientResponse> adviseStream(ChatClientRequest chatClientRequest, StreamAdvisorChain streamAdvisorChain) {

        beforeRequest(chatClientRequest);
        Flux<ChatClientResponse> chatClientResponseFlux = streamAdvisorChain.nextStream(chatClientRequest);


        return chatClientResponseFlux;
    }



    @Override
    public String getName() {
        return "WaveTokenCalculateAdvisor";
    }

    @Override
    public int getOrder() {
        return 0;
    }

}